//! Tool call response types.
//!
//! This module contains types for tool call responses:
//! - `InferenceResponseToolCall` - A validated tool call returned from inference
//! - `ToolCallChunk` - A streaming chunk of a tool call

#[cfg(feature = "pyo3")]
use pyo3::prelude::*;
use schemars::JsonSchema;
use serde::{Deserialize, Serialize, Serializer};
use serde_json::Value;
use tensorzero_derive::export_schema;

use super::config::ToolCallConfig;
use super::wire::ToolCall;

/// An InferenceResponseToolCall is a request by a model to call a Tool
/// in the form that we return to the client / ClickHouse
/// This includes some synactic sugar (parsing / validation of the tool arguments)
/// in the `arguments` field and the name in the `name` field.
/// We support looping this back through the TensorZero inference API via the ToolCallWrapper
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize, ts_rs::TS, JsonSchema)]
#[ts(export)]
#[cfg_attr(feature = "pyo3", pyclass(str))]
#[export_schema]
pub struct InferenceResponseToolCall {
    /// A Tool Call ID to match up with tool call responses. See #4058.
    pub id: String,

    /// The name of the tool to call, as generated by the model.
    pub raw_name: String,

    /// The raw arguments JSON string of the tool to call, as generated by the model.
    pub raw_arguments: String,

    /// The name of the tool to call, validated against tool configs. If not present, it means the tool call was invalid.
    pub name: Option<String>,

    /// The arguments of the tool to call, validated against tool configs. If not present, it means the tool call arguments were invalid.
    pub arguments: Option<Value>,
}

impl std::fmt::Display for InferenceResponseToolCall {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let json = serde_json::to_string_pretty(self).map_err(|_| std::fmt::Error)?;
        write!(f, "{json}")
    }
}

#[cfg(feature = "pyo3")]
#[pymethods]
impl InferenceResponseToolCall {
    pub fn __repr__(&self) -> String {
        self.to_string()
    }
}

impl InferenceResponseToolCall {
    /// Validates that a ToolCall is compliant with the ToolCallConfig
    /// First, it finds the ToolConfig for the ToolCall
    /// Then, it validates the ToolCall arguments against the ToolConfig
    pub async fn new(tool_call: ToolCall, tool_cfg: Option<&ToolCallConfig>) -> Self {
        // Check if this is a function tool
        let function_tool = tool_cfg.and_then(|t| t.get_function_tool(&tool_call.name));

        // Check if this is a custom tool
        let is_custom_tool = tool_cfg
            .map(|t| {
                t.openai_custom_tools
                    .iter()
                    .any(|ct| ct.name == tool_call.name)
            })
            .unwrap_or(false);

        // Set parsed_name if tool exists (either function or custom)
        let parsed_name = if function_tool.is_some() || is_custom_tool {
            Some(tool_call.name.clone())
        } else {
            None
        };

        // Validate arguments only for function tools (custom tools don't use JSON schemas)
        let parsed_arguments = if let Some(tool) = function_tool {
            if let Ok(arguments) = serde_json::from_str(&tool_call.arguments) {
                if tool.validate_arguments(&arguments).await.is_ok() {
                    Some(arguments)
                } else {
                    None
                }
            } else {
                None
            }
        } else {
            None
        };

        Self {
            arguments: parsed_arguments,
            id: tool_call.id,
            name: parsed_name,
            raw_arguments: tool_call.arguments.clone(),
            raw_name: tool_call.name.clone(),
        }
    }
}

#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
pub struct ToolCallChunk {
    pub id: String,
    #[serde(serialize_with = "serialize_option_string_as_empty")]
    pub raw_name: Option<String>,
    pub raw_arguments: String,
}

fn serialize_option_string_as_empty<S>(
    value: &Option<String>,
    serializer: S,
) -> Result<S::Ok, S::Error>
where
    S: Serializer,
{
    match value {
        Some(s) => serializer.serialize_str(s),
        None => serializer.serialize_str(""),
    }
}
